#include <gtest/gtest.h>

#include "kiva/dsl.h"

namespace kiva {
namespace {
using NodeHolder = std::unique_ptr<Node>;

TEST(DslNode, ImplicitConversionFromBasicTypes) {
  {
    DslNode n = false;
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kBool);
    EXPECT_EQ(nr->bool_value(), false);
  }
  {
    DslNode n = unsigned(0);
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kULong);
    EXPECT_EQ(nr->ulong_value(), 0);
  }
  {
    DslNode n = static_cast<unsigned long>(0);
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kULong);
    EXPECT_EQ(nr->ulong_value(), 0);
  }
  {
    DslNode n = static_cast<short>(1);
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kLong);
    EXPECT_EQ(nr->long_value(), 1);
  }
  {
    DslNode n = long(1);
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kLong);
    EXPECT_EQ(nr->long_value(), 1);
  }
  {
    DslNode n = 1;
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kLong);
    EXPECT_EQ(nr->long_value(), 1);
  }
  {
    DslNode n = 1.0;
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kFloat);
    EXPECT_FLOAT_EQ(nr->float_value(), 1);
  }
  {
    DslNode n = 1.0f;
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kFloat);
    EXPECT_FLOAT_EQ(nr->float_value(), 1);
  }
  {
    DslNode n = 1.0l;
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kFloat);
    EXPECT_FLOAT_EQ(nr->float_value(), 1);
  }
  {
    DslNode n = "string";
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kString);
    EXPECT_EQ(nr->string_value(), "string");
  }
  {
    DslNode n = std::string("string");
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kString);
    EXPECT_EQ(nr->string_value(), "string");
  }
  {
    DslNode n = std::string_view("string");
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kString);
    EXPECT_EQ(nr->string_value(), "string");
  }
}

TEST(DslNode, ConstructedFromInitializerList) {
  {
    DslNode n = {1, 2, 3};
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kTuple);
    EXPECT_EQ(nr->num_children(), 3);

    for (auto&& child : nr->children()) {
      EXPECT_EQ(child->type(), NodeType::kLong);
    }
  }
  {
    DslNode n = {};
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kTuple);
    EXPECT_TRUE(nr->children().empty());
  }
  {
    DslNode n = {"string", 1.0, false};
    auto nr = MakeNode(n);
    EXPECT_EQ(nr->type(), NodeType::kTuple);

    const NodeType expected_types[] = {
        NodeType::kString,
        NodeType::kFloat,
        NodeType::kBool,
    };

    for (int i = 0; i < nr->num_children(); ++i) {
      EXPECT_EQ(expected_types[i], nr->child_at(i)->type());
    }
  }
  {
    DslNode n = {
        {"foo", 10},
        {"bar", 2.0},
        {false, true},
    };
    auto nr = MakeNode(n);
    for (auto&& child : nr->children()) {
      EXPECT_EQ(child->type(), NodeType::kTuple);
      EXPECT_EQ(child->num_children(), 2);
    }
  }
}

TEST(DslNode, Subscript) {
  {
    auto v = MakeNode(DslNode(2)[2]);
    EXPECT_EQ(v->type(), NodeType::kSubscript);
    EXPECT_EQ(v->num_children(), 2);
    EXPECT_EQ(v->child_at(1)->type(), NodeType::kLong);
  }
}

TEST(DslNode, Call) {
  {
    DslNode n(1);
    auto v = MakeNode(n());
    EXPECT_EQ(v->type(), NodeType::kCall);
    EXPECT_EQ(v->num_children(), 2);
    EXPECT_EQ(v->child_at(1)->num_children(), 0);
    EXPECT_EQ(v->child_at(0)->type(), NodeType::kLong);
    EXPECT_EQ(v->child_at(1)->type(), NodeType::kTuple);
  }
  {
    DslNode n(1);
    auto v = MakeNode(n({false, 1.0}));
    EXPECT_EQ(v->type(), NodeType::kCall);
    EXPECT_EQ(v->num_children(), 2);

    EXPECT_EQ(v->child_at(1)->num_children(), 2);
    EXPECT_EQ(v->child_at(1)->child_at(0)->type(), NodeType::kBool);
    EXPECT_EQ(v->child_at(1)->child_at(1)->type(), NodeType::kFloat);
  }

}

TEST(DslNode, Assign) {
  {
    auto nr = MakeNode((DslNode(1) = 1));
    EXPECT_EQ(nr->type(), NodeType::kAssign);
    EXPECT_EQ(nr->child_at(0)->type(), NodeType::kLong);
    EXPECT_EQ(nr->child_at(1)->type(), NodeType::kLong);
  }
}

}  // namespace

namespace dsl {
namespace {

#define KIVA_TEST_BINOP(Name, Op)                          \
  {                                                        \
    auto node = MakeNode("n"_n Op 2);                      \
    EXPECT_EQ(node->type(), NodeType::k##Name);            \
    EXPECT_EQ(node->num_children(), 2);                    \
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kName); \
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kLong); \
  }                                                        \
  {                                                        \
    auto node = MakeNode(2 Op "n"_n);                      \
    EXPECT_EQ(node->type(), NodeType::k##Name);            \
    EXPECT_EQ(node->num_children(), 2);                    \
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kLong); \
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kName); \
  }

// clang-format off
TEST(DslNode, BoolOp) {
  KIVA_BOOLOP_LIST(KIVA_TEST_BINOP)
}

TEST(DslNode, BinOp) {
  KIVA_BINOP_LIST(KIVA_TEST_BINOP)
} 

// clang-format on
#undef KIVA_TEST_BINOP

#define KIVA_TEST_AUGASSIGN(Name, Op)                      \
  {                                                        \
    auto node = MakeNode("n"_n Op 2);                      \
    EXPECT_EQ(node->type(), NodeType::kAugAssign);         \
    EXPECT_EQ(node->num_children(), 3);                    \
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kName); \
    ASSERT_EQ(node->child_at(1)->type(), NodeType::kLong); \
    EXPECT_EQ(node->child_at(1)->long_value(),             \
              Node::LongType(NodeType::k##Name));          \
    EXPECT_EQ(node->child_at(2)->type(), NodeType::kLong); \
  }                                                        \
  {                                                        \
    auto node = MakeNode(2 Op "n"_n);                      \
    EXPECT_EQ(node->type(), NodeType::kAugAssign);         \
    EXPECT_EQ(node->num_children(), 3);                    \
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kLong); \
    ASSERT_EQ(node->child_at(1)->type(), NodeType::kLong); \
    EXPECT_EQ(node->child_at(1)->long_value(),             \
              Node::LongType(NodeType::k##Name));          \
    EXPECT_EQ(node->child_at(2)->type(), NodeType::kName); \
  }

TEST(DslNode, AugAssign) {
  KIVA_AUGASSIGN_LIST(KIVA_TEST_AUGASSIGN)
}
#undef KIVA_TEST_AUGASSIGN

TEST(DslNode, UnaryOp){
#define KIVA_TEST_UNARYOP(Name, Op)                        \
  {                                                        \
    auto node = MakeNode(Op "n"_n);                        \
    EXPECT_EQ(node->type(), NodeType::k##Name);            \
    ASSERT_EQ(node->num_children(), 1);                    \
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kName); \
  }
    KIVA_UNARYOP_LIST(KIVA_TEST_UNARYOP)
#undef KIVA_TEST_UNARYOP
}

TEST(DslNode, Compare){
#define KIVA_TEST_COMPARE(Name, Op)                        \
  {                                                        \
    auto node = MakeNode(1 Op "name"_n Op 2);              \
    EXPECT_EQ(node->type(), NodeType::kCompare);           \
    EXPECT_EQ(node->num_children(), 5);                    \
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kLong); \
    ASSERT_EQ(node->child_at(1)->type(), NodeType::kLong); \
    EXPECT_EQ(node->child_at(1)->long_value(),             \
              Node::LongType(NodeType::k##Name));          \
    EXPECT_EQ(node->child_at(2)->type(), NodeType::kName); \
    ASSERT_EQ(node->child_at(3)->type(), NodeType::kLong); \
    EXPECT_EQ(node->child_at(3)->long_value(),             \
              Node::LongType(NodeType::k##Name));          \
    EXPECT_EQ(node->child_at(4)->type(), NodeType::kLong); \
  }

    KIVA_COMPARE_LIST(KIVA_TEST_COMPARE)
#undef KIVA_TEST_COMPARE
}

TEST(NameDslNode, LiteralOperator) {
  auto node = MakeNode("name"_n);
  EXPECT_EQ(node->type(), NodeType::kName);
  EXPECT_EQ(node->string_value(), "name");
}

TEST(DslNodeBuilder, class_) {
  {
    auto node = MakeNode(class_("name", {}));
    EXPECT_EQ(node->type(), NodeType::kClassDef);
    ASSERT_EQ(node->num_children(), 2);

    auto c0 = node->child_at(0);
    EXPECT_EQ(c0->type(), NodeType::kName);
    EXPECT_EQ(c0->string_value(), "name");

    auto c1 = node->child_at(1);
    EXPECT_EQ(c1->type(), NodeType::kBody);
    EXPECT_EQ(c1->num_children(), 0);
  }
  {
    auto node = MakeNode(class_("name", {"base"}, {}));
    EXPECT_EQ(node->type(), NodeType::kClassDef);
    ASSERT_EQ(node->num_children(), 3);

    auto c0 = node->child_at(0);
    EXPECT_EQ(c0->type(), NodeType::kName);
    EXPECT_EQ(c0->string_value(), "name");

    auto c1 = node->child_at(1);
    EXPECT_EQ(c1->type(), NodeType::kTuple);
    EXPECT_EQ(c1->num_children(), 1);

    auto c2 = node->child_at(2);
    EXPECT_EQ(c2->type(), NodeType::kBody);
    EXPECT_EQ(c2->num_children(), 0);
  }
}

template <typename T>
struct GetNodeType;

// T is a Builder.
template <template <NodeType> class T, NodeType Type>
struct GetNodeType<T<Type>> : std::integral_constant<NodeType, Type> {};

template <typename T>
class BuilderTemplateTest : public ::testing::Test {
 protected:
  static NodeType node_type() {
    using U = std::remove_cv_t<T>;
    return GetNodeType<U>{};
  }
};

#define KIVA_MAKE_TYPE_TEST_SUITE(Name, ...)           \
  template <typename T>                                \
  class Name##Test : public BuilderTemplateTest<T> {}; \
  using Name##Types = ::testing::Types<__VA_ARGS__>;   \
  TYPED_TEST_SUITE(Name##Test, Name##Types)

KIVA_MAKE_TYPE_TEST_SUITE(OptionalValueBuilder,
                          decltype(return_),
                          decltype(yield));

KIVA_MAKE_TYPE_TEST_SUITE(ForBuilder, for_, async_for);

TYPED_TEST(OptionalValueBuilderTest, Basic) {
  auto return_ = TypeParam{};
  auto node_type = TestFixture::node_type();
  {
    auto node = MakeNode(return_);
    EXPECT_EQ(node->type(), node_type);
    ASSERT_EQ(node->num_children(), 0);
  }
  {
    auto node = MakeNode(return_(0));
    EXPECT_EQ(node->type(), node_type);
    ASSERT_EQ(node->num_children(), 1);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kLong);
  }
  {
    auto node = MakeNode(return_({1, 2}));
    EXPECT_EQ(node->type(), node_type);
    ASSERT_EQ(node->num_children(), 1);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kTuple);
  }
}

TYPED_TEST(ForBuilderTest, Basic) {
  using for_ = TypeParam;
  auto node = MakeNode(for_("i", "j", {}));
  EXPECT_EQ(node->type(), TestFixture::node_type());
  ASSERT_EQ(node->num_children(), 3);
  EXPECT_EQ(node->child_at(0)->type(), NodeType::kString);
  EXPECT_EQ(node->child_at(1)->type(), NodeType::kString);
  EXPECT_EQ(node->child_at(2)->type(), NodeType::kBody);
}

KIVA_MAKE_TYPE_TEST_SUITE(FunctionDefBuilder, def, async_def);

TYPED_TEST(FunctionDefBuilderTest, Basic) {
  using def = TypeParam;
  {
    auto node = MakeNode(def("name", {"arg1", "arg2"}, {}));
    EXPECT_EQ(node->type(), TestFixture::node_type());
    ASSERT_EQ(node->num_children(), 3);

    auto c0 = node->child_at(0);
    EXPECT_EQ(c0->type(), NodeType::kName);
    EXPECT_EQ(c0->string_value(), "name");

    auto c1 = node->child_at(1);
    EXPECT_EQ(c1->num_children(), 2);
    EXPECT_EQ(c1->type(), NodeType::kFormalArguments);

    auto c2 = node->child_at(2);
    EXPECT_EQ(c2->type(), NodeType::kBody);
    EXPECT_EQ(c2->num_children(), 0);
  }
  {
    auto node = MakeNode(def("name", {}, "int", {}));
    EXPECT_EQ(node->type(), TestFixture::node_type());
    ASSERT_EQ(node->num_children(), 4);
    
    auto c2 = node->child_at(2);
    EXPECT_EQ(c2->type(), NodeType::kString);
    EXPECT_EQ(c2->string_value(), "int");
  }
}

KIVA_MAKE_TYPE_TEST_SUITE(ConditionBodyBuilder, while_, if_, elif);

TYPED_TEST(ConditionBodyBuilderTest, Basic) {
  using while_ = TypeParam;
  auto node_type = TestFixture::node_type();
  {
    NodeHolder node = MakeNode(while_(false, {}));
    EXPECT_EQ(node->type(), node_type);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kBool);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kBody);
  }
}

KIVA_MAKE_TYPE_TEST_SUITE(BodyBuilder, else_, try_, finally, del);
TYPED_TEST(BodyBuilderTest, Basic) {
  using del = TypeParam;
  NodeType node_type = TestFixture::node_type();
  {
    auto node = MakeNode(del{1});
    EXPECT_EQ(node->type(), node_type);
    EXPECT_EQ(node->num_children(), 1);
  }
  {
    auto node = MakeNode(del{1, 2, 3});
    EXPECT_EQ(node->type(), node_type);
    EXPECT_EQ(node->num_children(), 3);
  }
}

KIVA_MAKE_TYPE_TEST_SUITE(WithBuilder, with, async_with);

TYPED_TEST(WithBuilderTest, Basic) {
  using with = TypeParam;
  NodeType node_type = TestFixture::node_type();
  {
    NodeHolder node = MakeNode(with("open"_n("path/file").as("f"), {}));
    EXPECT_EQ(node->type(), node_type);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kAlias);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kBody);
  }
}

TEST(DslNodeBuilder, RaiseBuilder) {
  {
    auto node = MakeNode(raise);
    EXPECT_EQ(node->type(), NodeType::kRaise);
  }
  {
    auto node = MakeNode(raise(1));
    EXPECT_EQ(node->type(), NodeType::kRaise);
    ASSERT_EQ(node->num_children(), 1);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kLong);
  }
  {
    auto node = MakeNode(raise(1, "from"));
    EXPECT_EQ(node->type(), NodeType::kRaise);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kLong);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kString);
  }
}

TEST(DslNodeBuilder, except) {
  {
    auto node = MakeNode(except{1, 2, 3, 4});
    EXPECT_EQ(node->type(), NodeType::kExcept);
    ASSERT_EQ(node->num_children(), 1);
    EXPECT_EQ(node->child_at(0)->num_children(), 4);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kBody);
  }
  {
    auto node = MakeNode(except("TypeError", {}));
    EXPECT_EQ(node->type(), NodeType::kExcept);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kString);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kBody);
  }
  {
    auto node = MakeNode(except("ValueError", "e", {}));
    ASSERT_EQ(node->num_children(), 3);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kString);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kName);
    EXPECT_EQ(node->child_at(2)->type(), NodeType::kBody);
  }
}

TEST(DslNodeBuilder, assert_) {
  {
    auto node = MakeNode(assert_(true));
    EXPECT_EQ(node->type(), NodeType::kAssert);
    ASSERT_EQ(node->num_children(), 1);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kBool);
  }
  {
    auto node = MakeNode(assert_(true, "msg"));
    EXPECT_EQ(node->type(), NodeType::kAssert);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kBool);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kString);
  }
}

KIVA_MAKE_TYPE_TEST_SUITE(NamesBuilder, import, global, nonlocal);

TYPED_TEST(NamesBuilderTest, Basic) {
  using import = TypeParam;
  auto node_type = TestFixture::node_type();
  {
    auto node = MakeNode(import{});
    EXPECT_EQ(node->type(), node_type);
    EXPECT_EQ(node->num_children(), 0);
  }
  {
    auto node = MakeNode(import{1});
    EXPECT_EQ(node->type(), node_type);
    EXPECT_EQ(node->num_children(), 1);
  }
  {
    auto node = MakeNode(import{1, 2, 3});
    EXPECT_EQ(node->type(), node_type);
    EXPECT_EQ(node->num_children(), 3);
  }
}

KIVA_MAKE_TYPE_TEST_SUITE(OddBallBuilder,
                          decltype(pass),
                          decltype(break_),
                          decltype(continue_));

TYPED_TEST(OddBallBuilderTest, Basic) {
  using OddBall = TypeParam;
  auto node_type = TestFixture::node_type();
  {
    NodeHolder node = MakeNode(OddBall{});
    EXPECT_EQ(node->type(), node_type);
    EXPECT_EQ(node->num_children(), 0);
  }
}

KIVA_MAKE_TYPE_TEST_SUITE(ValueBuilder, yield_from, await);

TYPED_TEST(ValueBuilderTest, Basic) {
  using yield_from = TypeParam;
  auto node_type = TestFixture::node_type();
  {
    auto node = MakeNode(yield_from{1});
    EXPECT_EQ(node->type(), node_type);
    ASSERT_EQ(node->num_children(), 1);
  }
}

TEST(DslNodeBuilder, from) {
  {
    auto node = MakeNode(from("os").import("path"));
    EXPECT_EQ(node->type(), NodeType::kImportFrom);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kName);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kString);
  }
  {
    auto node = MakeNode(from("re").import({"sub", "compile"}));
    EXPECT_EQ(node->type(), NodeType::kImportFrom);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kName);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kTuple);
  }
}

TEST(DslNodeBuilder, lambda) {
  {
    auto node = MakeNode(lambda({"a", "b"}, "a"_n + "b"_n));
    EXPECT_EQ(node->type(), NodeType::kLambda);
    ASSERT_EQ(node->num_children(), 2);
    EXPECT_EQ(node->child_at(0)->type(), NodeType::kTuple);
    EXPECT_EQ(node->child_at(1)->type(), NodeType::kAdd);
  }
}

}  // namespace
}  // namespace dsl
}  // namespace kiva
